import datetime
import multiprocessing
import psycopg2
import time
import argparse
import json
import redis
from kafka import KafkaProducer, KafkaConsumer

KAFKA_HOST = "localhost"
KAFKA_PORT = 9092
REDIS_HOST = 'localhost'
REDIS_PORT = 6379
PG_HOST = 'localhost'
PG_PORT = 5432
PG_DB = 'adc'
PG_USER = 'sqluser'
PG_PWD = "sqlpassword"
CAS_HOST = 'localhost'
CAS_TOKEN = ''
UPDATE_CYCLE = 60 * 60
IAM_AUTH_LOGS = "iam_auth_logs"
IAM_LOGIN_LOGS = "iam_login_logs"
CAS_LOGS = "cas"
AUTH3_HANDLER = "auth3_handler"
AUTH3_KEY = 'auth3'
kafka_producer = KafkaProducer(bootstrap_servers=KAFKA_HOST + ":" + str(KAFKA_PORT))
pg_connection = psycopg2.connect(database=PG_DB, user=PG_USER, password=PG_PWD, host=PG_HOST, port=PG_PORT)
redis_client = redis.StrictRedis(host='localhost', port=6379)


class StrategyDetection:

    def __init__(self, topic: str):
        # strategy_renew 更新策略时间参数
        self.cas_server = CAS_HOST
        self.x_token = CAS_TOKEN
        self.kafka_producer = KafkaProducer(bootstrap_servers=KAFKA_HOST + ":" + str(KAFKA_PORT))
        self.pg_connection = psycopg2.connect(database=PG_DB, user=PG_USER, password=PG_PWD,
                                              host=PG_HOST, port=PG_PORT)
        self.redis_client = redis.StrictRedis(host='localhost', port=6379)
        self.topic = topic
        self.last_update_timestamp = 0
        self.update_cycle = UPDATE_CYCLE
        self.stg_info = dict()
        self.manager = self.get_auth3_manager()  # dict
        self.cas_fail_list = ['STG/RES-2']
        self.cas_accumulate_uid_list = ['STG/RES-1', 'STG/RES-3']
        self._lockid = None
        # 判定的策略UID列表
        self.strategy = {
            # 以下为格式举例
            # "STG/RES-1": {
            #     "name": '资源请求频次过高',
            #     "desc": '资源请求频次达到<span style="color:red">100次/秒</span>，'
            #             '<span style="color:red">注销该用户登录状态</span>（强制退出），要求该用户登录。',
            #     "trigger_times": 100,
            #     "trigger_period": 1 * 1000,    # 单位为ms
            #     "mark": 0,  # 标记物
            #     "mark_time": None   # 初始标记时间
            # },
            # "STG/RES-3": {
            #     "name": '资源请求总量过高',
            #     "desc": '资源请求总量达到<span style="color:red">100000次/时</span>，'
            #             '<span style="color:red">告知管理员预警提示</span>。',
            #     "trigger_times": 100000,
            #     "trigger_period": 60 * 60 * 1000,  # 单位为ms
            #     "mark": 0,  # 标记物
            #     "mark_time": None   # 初始标记时间
            # },
        }

    def get_auth3_manager(self):
        res = self.redis_client.get(AUTH3_KEY)
        if not res:
            res = dict()
        else:
            res = json.loads(res)
        self.redis_client.set(AUTH3_KEY, '{}')
        return res

    def update_strategy(self):
        # 更新当前生效策略
        # uid_list = list(self.strategy.keys())
        print("更新当前策略表")
        self.last_update_timestamp = time.time()
        sql = "SELECT NAME, DESCRIPTION, UID, ENABLED, TRIGGER_TIMES, TRIGGER_PERIOD FROM ZTP.STRATEGY_STRATEGY" \
              " WHERE ENABLED = TRUE"
        with pg_connection.cursor() as cur:
            cur.execute(sql)
            rows = cur.fetchall()
            print("执行sql %s " % sql)
        # ToDo 考虑下渠道的区别
        print('Get %s line result' % (len(rows)))
        for row in rows:
            name = row[0]
            desc = row[1]
            uid = row[2]
            trigger_times = row[4]
            trigger_period = row[5]
            if not self.strategy.get(uid):
                self.strategy[uid] = dict()
            self.strategy[uid]["name"] = name
            self.strategy[uid]["desc"] = desc
            self.strategy[uid]["trigger_times"] = trigger_times
            self.strategy[uid]["trigger_period"] = trigger_period * 60 * 1000  # 分钟转换为ms
            self.strategy[uid]["mark"] = 0
            if uid == "STG/RES-2":
                # 越权访问, 触发即通知
                self.strategy[uid]["trigger_times"] = 0     # 一次触发
                self.strategy[uid]["trigger_period"] = 99999999999    # 99999999999 表示无周期策略,即累计策略
            if uid == "STG/RES-3":
                self.strategy[uid]["trigger_period"] = 99999999999  # 0表示无周期策略,即累计策略
            self.strategy[uid]["mark_time"] = None
            self.stg_info[uid] = dict()
        print("更新策略状态完毕 %s " % datetime.datetime.now())
        print('获取到策略集 %s' % self.strategy.keys())

    def detect_with_message(self, topic, message: dict):
        # message, 经过json.dumps转换的数据 message = json.loads(raw_message.value.decode('utf-8'))
        if topic == CAS_LOGS:
            print('接收 cas 鉴权日志')
            message_time = int(time.mktime(time.strptime(message.get("auth_time"), "%Y-%m-%d %H:%M:%S.%f")) * 1000) + \
                int(message.get("auth_time").split(".")[-1])
            # STG/RES-1 //STG/RES-3 // STG/RES-2 频次
            username = message.get("username")
            client_ip = message.get("client_ip")
            if username:
                # 更新STG/AUTH-3 判定
                self.update_auth3_strategy(username, message_time, client_ip, message.get("auth_time"))
                for uid in self.strategy.keys():
                    # if uid in self.cas_accumulate_uid_list:
                    #     self.handler_cycle_strategy(uid, username, message, client_ip, message_time)
                    # else:
                    self.handler_cycle_strategy2(uid, username, message, client_ip, message_time)
            # if message.get("auth_result") in ['VALIDATE_SUCCESS', "VALIDATE_FAIL", "AUTHORIZE_FAIL"] and username:
            #     for uid in self.strategy.keys():
            #         # 限定策略判定
            #         if uid not in self.cas_uid_list:
            #             continue
            #         self.handler_cycle_strategy(uid, username, message, client_ip, message_time)

        # 登录认证
        if topic == IAM_LOGIN_LOGS:
            print("接收 iam 登录日志")
            """STG/BASE-5 // STG/AUTH-13 // STG/AUTH-3 // STG/AUTH-5"""
            username = message.get("USER_NAME")
            client_ip = message.get("IP")
            message_time = int(time.mktime(time.strptime(message.get("CREATE_TIME"), "%Y-%m-%d %H:%M:%S")) * 1000)
            for uid in self.strategy.keys():
                if uid == "STG/BASE-5":
                    # 异地认证
                    pass
                if uid == "STG/AUTH-13":
                    # 认证失败冻结
                    self.handler_cycle_strategy2(uid, username, message, client_ip, message_time)
                if uid == "STG/AUTH-3":
                    # 更新闲时登出策略时间
                    self.update_auth3_strategy(username, message_time, client_ip, message.get('CREATE_TIME'))
                if uid == "STG/AUTH-5" and message.get("STATUS") == "成功":
                    # 异常时间登录
                    login_time_hour = time.strptime(message.get("CREATE_TIME"), "%Y-%m-%d %H:%M:%S").tm_hour
                    if login_time_hour >= 1 and login_time_hour <= 4:   # 1am - 5am
                        print("异常登录时间判定 %s %s" % (username, uid))
                        self.push_kafka(uid=uid, client_ip=client_ip, username=username,
                                        auth_time=message.get("CREATE_TIME"))
        # 鉴权认证
        if topic == IAM_AUTH_LOGS:
            print("接收 iam 鉴权日志")
            username = message.get("user_name")
            client_ip = message.get("IP")
            message_time = int(time.mktime(time.strptime(message.get("CREATE_TIME"), "%Y-%m-%d %H:%M:%S")) * 1000)
            for uid in self.strategy.keys():
                if uid == "STG/AUTH-5":
                    # 异常时间登录
                    login_time_hour = time.strptime(message.get("CREATE_TIME"), "%Y-%m-%d %H:%M:%S").tm_hour
                    if login_time_hour >= 1 and login_time_hour <= 4:
                        print("异常登录时间判定 %s %s" % (username, uid))
                        self.push_kafka(uid=uid, client_ip=client_ip, username=username,
                                        auth_time=message.get("CREATE_TIME"))
                if uid == "STG/AUTH-3":
                    # 更新闲时登出策略时间
                    self.update_auth3_strategy(username, message_time, client_ip, message.get('CREATE_TIME'))

    def update_auth3_strategy(self, username, message_time, client_ip=None, auth_time=None):
        # 更新STG/AUTH-3 闲时登出时间
        while True:
            locked = self.get_redis_lock()
            if not locked:
                time.sleep(1)
                continue
            else:
                break
        self.manager = self.get_auth3_manager()
        print('更新闲时登出 策略 stg/auth-3 %s ' % username)
        uid = "STG/AUTH-3"
        if not self.manager.get(username):
            self.manager[username] = dict()
        if message_time > self.manager[username].get('message_time', 0):
            user_dict = {
                "uid": uid, "message_time": message_time, "client_ip": client_ip,
                "auth_time": auth_time
            }
            self.manager[username] = user_dict
            # print(self.manager)
        self.redis_client.set(AUTH3_KEY, json.dumps(self.manager))
        self.release_redis_lock()

    def handler_auth3_strategy(self):
        # 更新STG/AUTH-3 闲时登出时间
        time.sleep(10)
        print("进入监控闲时登出策略....")
        while True:
            del_user = []
            time.sleep(1)
            self.manager = self.get_auth3_manager()
            for username in self.manager.keys():
                uid = self.manager[username]["uid"]
                client_ip = self.manager[username]['client_ip']
                auth_time = self.manager[username]['auth_time']
                now_time = time.time() * 1000
                if now_time - self.manager[username]["message_time"] >= 120 * 60 * 1000:
                    print("触发闲时登出策略 %s" % username)
                    self.push_kafka(uid, client_ip, username, auth_time)
                    # 重置
                    del_user.append(username)
                    # del self.manager[username]
            # 更新auth3 触发情况
            if del_user:
                while True:
                    locked = self.get_redis_lock()
                    if not locked:
                        time.sleep(1)
                        continue
                    else:
                        new_manager = self.get_auth3_manager()
                        for username in del_user:
                            if new_manager.get('username', {}).get('message_time') == \
                                    self.manager[username]['message_time']:
                                del new_manager[username]
                        self.redis_client.set(AUTH3_KEY, json.dumps(new_manager))
                        self.release_redis_lock()
                        break

    def cas_server_callback(self):
        # todo 调用cas专用
        pass

    def publish_mqtt(self, message):
        # ToDo
        pass

    def run(self):
        # 监听kafka队列, 转换消息
        self.update_strategy()
        if self.topic == AUTH3_HANDLER:
            return self.handler_auth3_strategy()
        kafka_consumer = get_kafka_consumer(self.topic)
        print("开始监听kafka队列数据......")
        for msg in kafka_consumer:
            print("Get msg %s " % str(msg.value))
            now = time.time()
            if now - self.last_update_timestamp >= self.update_cycle:
                # 更新当前策略情况
                self.update_strategy()
            self.detect_with_message(self.topic, msg.value)

    def handler_cycle_strategy2(self, uid, username, message, client_ip, message_time):
        # 处理连续时间任务
        if not self.stg_info[uid].get(username):
            # print(self.stg_info[uid][username])
            self.stg_info[uid][username] = list()   # 队列处理连续策略逻辑
        if uid in self.cas_fail_list:
            print("处理失败任务策略 %s " % uid)
            if message.get("auth_result") in ["VALIDATE_FAIL", "AUTHORIZE_FAIL"]:
                # 推送积累数据
                self.stg_info[uid][username].append(
                    {"message_time": message_time,
                     "uid": uid, "client_ip": client_ip}
                )
                # print(self.stg_info[uid][username])
            elif message.get("auth_result") in ['VALIDATE_SUCCESS']:
                # 清空数组
                self.stg_info[uid][username] = list()
        elif uid in self.cas_accumulate_uid_list:
            print("处理连续积累性任务策略 %s " % uid)
            if message.get("auth_result") in ["VALIDATE_FAIL", "AUTHORIZE_FAIL", "VALIDATE_SUCCESS"]:
                # 推送积累数据
                self.stg_info[uid][username].append(
                    {"message_time": message_time,
                     "uid": uid, "client_ip": client_ip}
                )
        elif uid == "STG/AUTH-13":
            print("处理连续性任务策略 %s " % uid)
            # 认证失败冻结(博威IAM)
            # print(message.get("STATUS"))
            if message.get("STATUS") == "失败":
                # 推送积累数据
                self.stg_info[uid][username].append(
                    {"message_time": message_time,
                     "uid": uid, "client_ip": client_ip}
                )
            elif message.get("STATUS") == "成功":
                # 清空数组
                # print("重置清空数据")
                self.stg_info[uid][username] = list()
            # todo 以下为cas logo支持
            if message.get("auth_result") in ['AUTH_FAIL']:
                pass
            elif message.get("auth_result") in ['AUTH_SUCCESS']:
                pass
        else:
            return
        # 检测是否触发策略
        if (len(self.stg_info[uid][username]) < self.strategy[uid]['trigger_times']) or not self.stg_info[uid][username]:
            print("未触发策略 % s" % uid)
            return
        # elif len(self.stg_info[uid][username]) == self.strategy[uid]['trigger_times']:
        else:
            if self.stg_info[uid][username][-1]["message_time"] - self.stg_info[uid][username][0]['message_time'] \
                    <= self.strategy[uid]['trigger_period']:
                # 触发策略,回调cas
                self.cas_server_callback()
                self.publish_mqtt(message)
                self.push_kafka(uid=uid, client_ip=client_ip, username=username)
                # 清空数组
                self.stg_info[uid][username] = list()
            else:
                # 推出
                print("未满足策略触发条件 %s" % uid)
                if self.stg_info[uid][username]:
                    self.stg_info[uid][username].pop(0)

    def get_redis_lock(self, name="AUTH3_LOCK", lock_time=20):
        val = self.redis_client.get(name)
        if val:
            return False
        self._lockid = time.time()
        val_lock = 'locked %s' % self._lockid
        print(val_lock)
        re = self.redis_client.set(name, val_lock, ex=lock_time, nx=True)
        return re

    def release_redis_lock(self, name='AUTH3_LOCK'):
        val = self.redis_client.get(name)
        if not val:
            return True
        if str(val, 'utf-8') == "locked %s" % self._lockid:
            self.redis_client.delete(name)
            print("release lock %s" % self._lockid)
            return True
        return False


    # def handler_cycle_strategy(self, uid, username, message, client_ip, message_time):
    #     print("处理连续积累任务策略 %s " % uid)
    #     if message.get("auth_result") not in ["VALIDATE_FAIL", "AUTHORIZE_FAIL", 'VALIDATE_SUCCESS']:
    #         return
    #     print(self.stg_info[uid].get(username))
    #     if not self.stg_info[uid].get(username):
    #         self.stg_info[uid][username] = dict()   # map用于处理累计策略逻辑
    #     self.stg_info[uid][username] = self.strategy[uid]
    #     if not self.stg_info[uid][username]['mark_time']:
    #         self.stg_info[uid][username]['mark_time'] = message_time
    #         self.stg_info[uid][username]["mark"] = 1
    #     else:
    #         self.stg_info[uid][username]["mark"] += 1
    #     # 判定是否在策略触发周期内
    #     if self.stg_info[uid][username]['trigger_period'] and message_time - self.stg_info[uid][username]['mark_time'] >= \
    #             self.stg_info[uid][username]['trigger_period']:
    #         # 超期,重置
    #         self.stg_info[uid][username]['mark_time'] = message_time
    #         self.stg_info[uid][username]["mark"] = 1
    #     else:
    #         if self.stg_info[uid][username]['mark'] >= self.stg_info[uid][username]['trigger_times']:
    #             # 触发策略,回调cas
    #             self.cas_server_callback()
    #             self.publish_mqtt(message)
    #             self.push_kafka(uid=uid, client_ip=client_ip, username=username)
    #             self.stg_info[uid][username]['mark_time'] = None
    #             self.stg_info[uid][username]["mark"] = 0


    def get_strategy(self):
        return self.strategy

    def push_kafka(self, uid, client_ip, username, auth_time=None):
        print('--------推送策略触发日志 %s %s--------' % (uid, username))
        strategy_data = {
            "uid": uid,
            "auth_type": None,
            "auth_time": auth_time,
            "client_ip": client_ip,
            "username": username,
            "auth_result": None,
            "url": None
        }
        self.kafka_producer.send("strategy", json.dumps(strategy_data).encode("utf-8"))


def get_kafka_consumer(topic,  k_host=KAFKA_HOST, k_port=KAFKA_PORT, user=None,
                       pwd=None, auto_commit=True, group_id="strategy"):
    if user and pwd:
        consumer = KafkaConsumer(
            topic,
            sasl_mechanism="PLAIN",
            security_protocol="SASL_PLAINTEXT",
            sasl_plain_username=user,
            sasl_plain_password=pwd,
            bootstrap_servers=f"{k_host}:{k_port}",
            value_deserializer=lambda m: json.loads(m.decode('unicode_escape')),
            group_id=group_id, auto_offset_reset="earliest",
            enable_auto_commit=auto_commit,
            session_timeout_ms=30000,
            max_poll_interval_ms=300000 * 12 * 6
        )
    else:
        consumer = KafkaConsumer(
            topic,
            bootstrap_servers=f"{k_host}:{k_port}",
            value_deserializer=lambda m: json.loads(m.decode('unicode_escape')),
            group_id=group_id, auto_offset_reset="earliest",
            enable_auto_commit=auto_commit,
            session_timeout_ms=300000,
            max_poll_interval_ms=300000 * 12 * 6
        )
    return consumer


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--topic', type=str,
        required=True,
        help='Kafka Topic名称'
    )
    args = parser.parse_args()
    if args.topic == IAM_LOGIN_LOGS:
        iam_login_consumer = StrategyDetection(IAM_LOGIN_LOGS)
        iam_login_consumer.run()
    elif args.topic == IAM_AUTH_LOGS:
        iam_auth_consumer = StrategyDetection(IAM_AUTH_LOGS)
        iam_auth_consumer.run()
    elif args.topic == CAS_LOGS:
        cas_consumer = StrategyDetection(CAS_LOGS)
        cas_consumer.run()
    elif args.topic == AUTH3_HANDLER:
        auth3_handler = StrategyDetection(AUTH3_HANDLER)
        auth3_handler.run()
    else:
        test_handler = StrategyDetection('test_handler')
        locked = test_handler.get_redis_lock()
        if locked:
            print(1111)
            time.sleep(10)
            test_handler.release_redis_lock()
        else:
            print('No locked result')


